import pandas as pd
import numpy as np
from datetime import datetime
import re


def clean_data(df):
    """
    通用数据清洗函数

    参数:
    df (pd.DataFrame): 输入的脏数据DataFrame

    返回:
    pd.DataFrame: 清洗后的数据
    """
    # 1. 数据备份
    cleaned_df = df.copy()

    # 2. 初步诊断报告
    original_rows = cleaned_df.shape[0]
    original_columns = cleaned_df.shape[1]
    print(f"原始数据: {original_rows} 行, {original_columns} 列")
    print(f"初始缺失值: {cleaned_df.isnull().sum().sum()} 个")

    # 3. 处理缺失值
    # 自动识别列类型
    num_cols = cleaned_df.select_dtypes(include=np.number).columns.tolist()
    cat_cols = cleaned_df.select_dtypes(include=['object', 'category']).columns.tolist()
    date_cols = [col for col in cleaned_df.columns if cleaned_df[col].dtype == 'datetime64[ns]']

    # 数值列填充
    for col in num_cols:
        if cleaned_df[col].isnull().sum() > 0:
            # 小比例缺失用中位数，大比例缺失用插值
            if cleaned_df[col].isnull().mean() < 0.3:
                cleaned_df[col].fillna(cleaned_df[col].median(), inplace=True)
            else:
                cleaned_df[col] = cleaned_df[col].interpolate()

    # 分类列填充
    for col in cat_cols:
        if cleaned_df[col].isnull().sum() > 0:
            mode_val = cleaned_df[col].mode()
            cleaned_df[col].fillna(mode_val[0] if not mode_val.empty else "Unknown", inplace=True)

    # 4. 处理异常值 (仅数值列)
    for col in num_cols:
        # 跳过二进制/布尔值列
        if cleaned_df[col].nunique() <= 2:
            continue

        # 稳健的异常值检测
        q1 = cleaned_df[col].quantile(0.25)
        q3 = cleaned_df[col].quantile(0.75)
        iqr = q3 - q1

        # 仅当IQR不为零时处理
        if iqr > 0:
            lower_bound = q1 - 1.5 * iqr
            upper_bound = q3 + 1.5 * iqr

            # 温和的缩尾处理，而非硬截断
            cleaned_df[col] = np.where(
                cleaned_df[col] < lower_bound,
                np.percentile(cleaned_df[col], 5),
                cleaned_df[col]
            )
            cleaned_df[col] = np.where(
                cleaned_df[col] > upper_bound,
                np.percentile(cleaned_df[col], 95),
                cleaned_df[col]
            )

    # 5. 处理重复值 (保留第一个出现的)
    cleaned_df.drop_duplicates(inplace=True)

    # 6. 数据类型修复
    # 尝试转换可能的日期列
    date_pattern_cols = [col for col in cleaned_df.columns if re.search(r'date|time|timestamp', col, re.I)]
    for col in date_pattern_cols + date_cols:
        try:
            cleaned_df[col] = pd.to_datetime(cleaned_df[col], errors='coerce')
        except:
            pass

    # 尝试转换数值列
    for col in cleaned_df.columns:
        if col not in num_cols and col not in date_cols and col not in date_pattern_cols:
            try:
                # 尝试转换数值
                converted = pd.to_numeric(cleaned_df[col], errors='raise')
                cleaned_df[col] = converted
                num_cols.append(col)
                if col in cat_cols:
                    cat_cols.remove(col)
            except:
                # 如果转换失败，尝试清理字符串
                if col in cat_cols:
                    cleaned_df[col] = cleaned_df[col].astype(str).str.strip()

    # 7. 文本数据标准化
    for col in cat_cols:
        # 统一大小写并去除首尾空格
        cleaned_df[col] = cleaned_df[col].astype(str).str.strip().str.title()

        # 替换常见错误表示
        cleaned_df[col] = cleaned_df[col].replace({
            r'\bNaN\b': 'Unknown',
            r'\bNull\b': 'Unknown',
            r'\bN/A\b': 'Unknown',
            r'\bNone\b': 'Unknown'
        }, regex=True)

    # 8. 特殊字符处理
    for col in cleaned_df.columns:
        if col in cat_cols:
            # 移除不可见字符和非ASCII字符
            cleaned_df[col] = cleaned_df[col].apply(
                lambda x: re.sub(r'[\x00-\x1F\x7F-\xFF]', '', x) if isinstance(x, str) else x
            )

    # 9. 结果报告
    final_rows = cleaned_df.shape[0]
    final_columns = cleaned_df.shape[1]
    removed_rows = original_rows - final_rows
    removed_missing = original_columns * original_rows - cleaned_df.count().sum()

    print("\n=== 清洗报告 ===")
    print(f"- 移除重复行: {removed_rows} 行")
    print(f"- 处理缺失值: {removed_missing} 个")
    print(f"清洗后数据: {final_rows} 行, {final_columns} 列")
    print(f"剩余缺失值: {cleaned_df.isnull().sum().sum()} 个")

    return cleaned_df


# ===== 使用示例 =====
if __name__ == "__main__":
    # 示例1: 创建脏数据测试
    dirty_data = pd.DataFrame({
        'user_id': [1, 2, 3, 4, 5, 5, 5],
        'name': ['John ', ' alice', 'BOB', None, 'NaN', '  David  ', 'Eva'],
        'age': [25, 30, -5, 40, None, 150, 35],
        'salary': [50000, 60000, 1000000, 45000, 70000, 30000, 75000],
        'join_date': ['2020-01-01', '2021-02-30', 'invalid', '2022-03-15', None, '2023-01-01', '2023-01-01'],
        'email': ['john@example.com', 'alice@', 'invalid', 'bob@example.com', None, 'david@example', 'eva@example.com']
    })

    print("===== 清洗前数据 =====")
    print(dirty_data)

    # 清洗数据
    cleaned_data = clean_data(dirty_data)

    print("\n===== 清洗后数据 =====")
    print(cleaned_data)

    # 示例2: 从CSV文件读取
    # cleaned_data = clean_data(pd.read_csv('dirty_data.csv'))
    # cleaned_data.to_csv('cleaned_data.csv', index=False)